from core.model import SimpleABMModel
import matplotlib.pyplot as plt
import numpy as np
import json
import os
from datetime import datetime

def load_config(config_path="config.json"):
    """加载配置文件"""
    with open(config_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def create_output_dir():
    """创建输出目录"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f"output_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)
    return output_dir

def plot_results(model, output_dir):
    """绘制结果"""
    # 获取所有代理的状态
    states = model.get_agent_states()
    
    # 提取意见和自信度
    opinions = [state['opinion'] for state in states]
    confidences = [state['confidence'] for state in states]
    
    # 绘制意见分布
    plt.figure(figsize=(10, 5))
    plt.hist(opinions, bins=20, alpha=0.7)
    plt.title('Opinion Distribution')
    plt.xlabel('Opinion')
    plt.ylabel('Count')
    plt.savefig(os.path.join(output_dir, 'opinion_distribution.png'))
    plt.close()
    
    # 绘制自信度分布
    plt.figure(figsize=(10, 5))
    plt.hist(confidences, bins=20, alpha=0.7)
    plt.title('Confidence Distribution')
    plt.xlabel('Confidence')
    plt.ylabel('Count')
    plt.savefig(os.path.join(output_dir, 'confidence_distribution.png'))
    plt.close()
    
    # 绘制意见-自信度散点图
    plt.figure(figsize=(10, 5))
    plt.scatter(opinions, confidences, alpha=0.5)
    plt.title('Opinion vs Confidence')
    plt.xlabel('Opinion')
    plt.ylabel('Confidence')
    plt.savefig(os.path.join(output_dir, 'opinion_vs_confidence.png'))
    plt.close()
    
    # 保存最终状态
    with open(os.path.join(output_dir, 'final_states.json'), 'w', encoding='utf-8') as f:
        json.dump(states, f, ensure_ascii=False, indent=4)

def main():
    """主函数"""
    # 创建输出目录
    output_dir = create_output_dir()
    
    # 加载配置
    config = load_config()
    
    # 创建模型
    model = SimpleABMModel(config_path="config.json")
    
    # 运行模型
    n_steps = config.get('n_steps', 100)
    for i in range(n_steps):
        model.step()
        if (i + 1) % 10 == 0:
            print(f"Step {i + 1}/{n_steps}")
    
    # 绘制结果
    plot_results(model, output_dir)
    print(f"Results saved to {output_dir}")

if __name__ == "__main__":
    main() 